import torchaudio
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write

# 1. 读取原始 MP3
melody, sr = torchaudio.load(r"C:\Users\33720\Desktop\好多好多日语句子\好多好多日语句子cmp3\5-生活长单词1-1.wav")

# 2. 加载 MusicGen 模型
model = MusicGen.get_pretrained("facebook/musicgen-melody")
model.set_generation_params(duration=16)  # 生成16秒

# 3. 文本提示：雷鬼风格
descriptions = ["reggae style, upbeat rhythm, one drop drum, skank guitar, lively bass"]

# 4. 生成雷鬼风格音乐
# expand 旋律到 len(descriptions)
wav = model.generate_with_chroma(descriptions, melody[None].expand(len(descriptions), -1, -1), sr)

# 5. 保存生成音频
for idx, one_wav in enumerate(wav):
    audio_write(f"reggae_{idx}", one_wav.cpu(), model.sample_rate, strategy="loudness")
